#!/usr/bin/env python
# -*- coding: ISO-8859-1 -*-

__author__  = 'bryan conrad'
__license__ = 'GPL'
__doc__ = """
This program is a Python version of the JointTrack renderer    
"""
import vtk
import os
import sys
from enthought.tvtk.api import tvtk

# FOR DEBUGGING
import foo

class pyJTRederer():
    """This class perfroms the main rendering duties for the image registration tool.
    Most methods from VTK are traited (enthought.tvtk.api) so that they can easily be configured using the handy configureTraits() call"""

    def __init__(self):
        
#---Options---#000000#FFFFFF----------------------------------------------------
        self.nIter = 50
        self.savePng = 0
        self.PrintFrameRate = 1 #print the frame rate to console during rotate
        #TODO: Check platform and choose the path automatically
        if sys.platform == 'linux2':
            self.fileNameVol = '/home/bryan/Data/vert1.vtk'  #Linux
        else:
            if os.path.exists('C:/Lib/jointtrack/branches/TestData/vert1.vtk'):
                self.fileNameVol = 'C:/Lib/jointtrack/branches/TestData/vert1.vtk' #Vista
            else:
                self.fileNameVol = 'C:/Data/2D3D/vert1/CT/vert1.vtk' #Windows
        self.fileNameImg = '.tif'
        self.fileNameOut = 'out.png'
        #fileName = 'C:/Bryan/Python/foot.vtk'
        self.mapper = 1 # 1 means you want to use texture mapping (if your hardware supports it)
        self.osr = 0 # Set offscreen rendering
        self.waitTime = 0   #How long should we wait between iterations of "rotate".
#-------------------------------------------------------------------------------

        
        #Setup basic rendering components.
        #For the renderer to work with the WX renderwindow, it has to be a vtk renderer, not the tvtk version
        # this does not seem to make sense to be because they both look like they return a valid VTK object (?)
##        self.vtkRen = vtk.vtkRenderer()
        self.ren = tvtk.Renderer()
        self.renWin = tvtk.RenderWindow()
        self.renWin.add_renderer(self.ren)
##        self.renWin.add_renderer(self.vtkRen)

        self.iren = tvtk.RenderWindowInteractor()
        self.style = tvtk.InteractorStyleTrackballCamera()
        # Other Style options
        #style = vtk.vtkInteractorStyleTrackbal()
        #style = vtk.vtkInteractorStyleJoystickCamera()
        #style = vtk.vtkInteractorStyleRubberBandZoom()
        self.renWin.size = 512, 512
        self.renWin.window_name = 'Python JointTrack'
        self.iren._set_render_window(self.renWin)
        
        #Create a VTK volume for the 3D CT volume
        self.volume = tvtk.Volume()
        
        # Create the reader for the volume data
        self.volumeReader = tvtk.StructuredPointsReader()
        self.volumeReader.file_name = self.fileNameVol
        
        # Create the reader for the fluoro data
        #TODO: setup a more general image reader so that other file formats can be handled (right now only TIFF)!
        self.imgReader = tvtk.TIFFReader()
        
        # Create transfer mapping scalar value to opacity
        self.opacityTransferFunction = tvtk.PiecewiseFunction()
        #Default settings for the opacity transfer funciton, these can be changed using SetTransferFunction
        self.opacityTransferFunction.add_point(0.0, 0.0)
        self.opacityTransferFunction.add_point(70, 0.0)
        self.opacityTransferFunction.add_point(255, 0.5)
        
        # Create transfer mapping scalar value to color
        self.colorTransferFunction = tvtk.ColorTransferFunction()
        #Default settings for the color transfer funciton, these can be changed using SetTransferFunction
        self.colorTransferFunction.add_rgb_point(0.0, 0.0, 0.0, 0.0)
        self.colorTransferFunction.add_rgb_point(72, 1.0, 1.0, 1.0)
        self.colorTransferFunction.add_rgb_point(255, 1.0, 1.0, 1.0)
        
        # The property describes how the data will look
        self.volumeProperty = tvtk.VolumeProperty()
        # Setup default properties
        self.volumeProperty.set_color(self.colorTransferFunction)
        self.volumeProperty.set_scalar_opacity(self.opacityTransferFunction)
        self.volumeProperty.shade = 1
        self.volumeProperty.interpolation_type = 'linear'

        #Set up the background (fluoro) image
        self.backgroundMapper = tvtk.ImageMapper()
        self.backgroundMapper.color_window = 256
        self.backgroundMapper.color_level = 127.5
        
        self.backgroundProp = tvtk.Actor2D()
        self.backgroundProp._set_mapper(self.backgroundMapper)
        self.backgroundPropProperties = self.backgroundProp._get_property()
        self.backgroundPropProperties.display_location = "background" 
        
        #Determine if texture mapping rendering is supported by checking with the texture mapper class.
        #Texture mapper is fast, but not supported on all graphics cards.
        #To see if it is supported, we setup texture mapping, and tthen check the property "isRenderSupported"
        #if it is not, then we revert to raycast mapper (although, i think texture mapper will automatically fall back to a supported
        # mapper, i am not sure which one)  
        self.volumeMapper = tvtk.VolumeTextureMapper3D()
        self.volumeMapper.set_input(self.volumeReader.get_output())
        
        self.volume._set_property(self.volumeProperty)
        self.volume._set_mapper(self.volumeMapper)
        
        #self.volumeReader.update()
        
        self._renderingSupported = self.volumeMapper.is_render_supported(self.volumeProperty)
        
        if self._renderingSupported:
            print "Texture Mapper will be used for rendering [fast]"
        else:
            print "Your graphics card does not support texture mapping.  Raycast Maper will be used for rendering [slow]"
            self.volumeMapper = tvtk.VolumeRayCastMapper()
            self.compositeFunction = tvtk.VolumeRayCastCompositeFunction()
            self.volumeMapper._set_volume_ray_cast_function(self.compositeFunction)
            self.volumeMapper.set_input(self.volumeReader.get_output())
        
##        self.vtkRen.AddVolume(self.volume)
        self.ren.add_volume(self.volume)
        
        #Setup a Window2Image filter so that the output of the 3D transformation can beused in further image processing (registration)
        self.w2if = tvtk.WindowToImageFilter()
        self.w2if._set_input(self.renWin)
        
        #Create image file writer
        self.PNGimageWriter = tvtk.PNGWriter()
        self.PNGimageWriter._set_input(self.w2if._get_output())
        #set default output iimage name as out.png
        self.PNGimageWriter.file_name = self.fileNameOut
        
        #The following code sets up a text mapper, to render text in the render window.
        #This is not a good idea for images that will be used for registration, but for output 
        # or debugging it could be useful.
        self.textMapper = tvtk.TextMapper()
        self.TextProperties = self.textMapper._get_text_property()
        self.TextProperties.font_family = 'arial'
        self.TextProperties.font_size = 20
        self.TextProperties.shadow = 1
        self.textActor = tvtk.Actor2D()
        self.textActor._set_mapper(self.textMapper)
        #TODO: the position property should probably be scaled to the window size
        self.textActor.position = (10,10)
##        self.vtkRen.AddActor2d(self.textActor)
        
        #for now the transform is not used, the camera is rotated instead.
        #TODO: rotate volume using transfrom
        self.transform = tvtk.Transform()
        self.transform._set_input(self.volume._get_user_transform())
        #Make sure that when the transform is used, that the correct rotation center is specified.
        self.origin = self.volume._get_center()
        
    def PrepareFrame(self):    
        """This function will prepare the render main window for rendering"""
        #What actually needs to be done here???
        #self.Set3DMapper()
        self.renWin.render()

    def Render(self):
        """This fucntion will render the superimposed 3D and 2D images."""
        print 'Beginning render'
##        self.vtkRen.ResetCamera()        
        self.ren.reset_camera()
        self.renWin.render()        
        
        def CheckAbort(obj, event):
            if obj.GetEventPending() != 0:
                obj.SetAbortRender(1)
 
        self.renWin.add_observer("AbortCheckEvent", CheckAbort)
        
        print self.volumeMapper
        print self.renWin
        
        self.iren.initialize()
        self.renWin.render()
        self.iren.start()
    
    def RotateVolume(self):
        """This function is only for demonstration/debugging.  It provide a good test for rendering speed."""    
        for i in range(0,self.nIter):
            self.renWin.render()
            time.sleep(self.waitTime)
            self.vtkRen.GetActiveCamera().Azimuth( 360/self.nIter )
            
            #Print framerate to console
            if self.PrintFrameRate:
                print '%.1f fps' % (1/self.vtkRen.GetLastRenderTimeInSeconds())
            
            # if (ren.GetLastRenderTimeInSeconds()) < (firstRender/10):
                # print 'Render error at iteration: ' + str(i) + ' of ' + str(nIter)
                # break        
            
            if self.savePng:
                outName = 'output/example' + str(i) + '.png'
                self.PNGimageWriter.fileName = outName
                self.w2if.Modified()
                self.PNGimageWriter.write()
                #textMapper2.SetInput('Saving projection to file: ' + outName)
            else:
                pass
                #textMapper2.SetInput('Projection not saved')
                
        print "Done"
        #t2=time.time()
        #print "time to render and save " + str(i) + " images: " + str(t2-t1) + "s"
        #x=raw_input()
        #renWin.HideCursor()
        
    def SetVolume(self, filename):
        """Set the path the the 3D volume"""
        #TODO: allow the program to handle more than one volume
        self.fileNameVol = filename
    
    def GetVolume(self):
        """Get the path the the 3D volume"""
        #TODO: allow the program to handle more than one volume        
        return self.fileNameVol
    
    def Set2DImage(self, filename):
        """Set the path the the 2D image"""
        self.fileNameImg = filename
    
    def Get2DImage(self):
        """Get the path the the 2D image"""
        return self.fileNameImg
    
    def SetOpacityTransferFunction(self,point1,point2,point3):
        """Create custom color transfer function (LUT) for the volume renderer.  Input 3 points as 2 element tuples (grayscale value, opacity)"""
        self.opacityTransferFunction.add_point(point1[0],point1[1])
        self.opacityTransferFunction.add_point(point2[0],point2[1])
        self.opacityTransferFunction.add_point(point3[0],point3[1])
    
    def SetColorTransferFunction(self,point1,point2,point3):
        """Create custom opacity transfer function (LUT) for the volume renderer.  Input 3 points as 2 element tuples (grayscale value, opacity)"""
        self.opacityTransferFunction.add_rgb_point(point1[0],point1[1],point1[2],point1[3])
        self.opacityTransferFunction.add_rgb_point(point2[0],point2[1],point2[2],point2[3])
        self.opacityTransferFunction.add_rgb_point(point3[0],point3[1],point3[2],point3[3])
    
    def SetOutputFileName(self,filename):
        """Specify the filename for the output image"""
        self.PNGimageWriter.file_name = filename
        
    def test(self):
        print "this is a test", self.nIter
        print "Now calling foo.bar"
        f=foo()
        f.bar
    def WriteText(self,inputText):
        self.textMapper._set_input(inputText)
    
    def UseCone(self):
        """For Debugging, use a cone source instead of rendering a volume"""            
        print 'Using Cone'
        cone = tvtk.ConeSource()
        cone.height = 3.0
        cone.radius = 1.0
        cone.resolution = 10
        coneMapper = tvtk.PolyDataMapper(input=cone.output)
        self.coneActor = tvtk.Actor(mapper=coneMapper)
    
    def UseImage(self):
        """For Debugging, use a cone source instead of rendering a volume"""            
        print 'Using Image'
        imageReader = tvtk.PNGReader(file_name='/home/bryan/out1.png')
        self.imageActor = tvtk.ImageActor(input=imageReader.output)

        
if __name__ == '__main__':
    print "not implemented"
    
